#include <fcntl.h>
#include <stdio.h>
#include <cmath>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <sndfile.h>
#include <set>
#include <vector>
#include <algorithm>
#include <map>
#include <boost/property_map.hpp>
#include <boost/pending/relaxed_heap.hpp>

using namespace std;

class Cmp {
 public:
  bool operator()(float *a, float *b) const {
    return (*a) > (*b);
  }
};

class ptr_map
{
 public:
  typedef float* key_type; 
  typedef int value_type;
  typedef boost::lvalue_property_map_tag category;
  typedef float*& reference;
  
  ptr_map(float *b, int *v) : base(b), val(v) {
    
  };

  ~ptr_map() {
    //delete[] val;
  }
  float *base;
  int *val;
};


int get(const ptr_map& p, float *key) {
  return p.val[key - p.base];
}

void put(const ptr_map& p, float *key, int value) {
  p.val[key - p.base] = value;
}

#define PI 3.14159
#define NUMK 32
#define ERB(x) (24.7*(4.37*(x)/1000.0 + 1.0))
#define EARQ 9.26449
#define MINBW 24.7
#define ORDER 1

set<int> neighborhoods[NUMK];
typedef boost::relaxed_heap<float*, Cmp, ptr_map> data_heap;
data_heap *sig_heaps[NUMK];

float *orig_sig;
float *sig;
float *prev_sig;
int sig_len;
float curr_coeffs[NUMK];

float *gen_gammatone(float f, int *num_samples) {
  double n = 3.0;
  double b = 1.019*ERB(f);

  *num_samples = (int)((80000.0)/(b));

  float *result = (float *) malloc((*num_samples)*sizeof(float));

  int x;
  double total = 0.0;
  for (x = 0; x < *num_samples; x++) {
    double t = x/44100.0;
    result[x] = pow(t,n)*exp(-2*PI*b*t)*cos(2*PI*f*t);
    total += pow(result[x], 2);
  }
  total = sqrt(total);
  for (x = 0; x < *num_samples; x++) {
    result[x] /= total;
    //result[x+1] /= total;
  }

  return result;
}

void compute_gammatones(float **gts, int *gt_lens, int num,
			float lf, float hf) {
  float stepfactor = EARQ * (log(hf + EARQ * MINBW) - 
			     log(lf + EARQ * MINBW)) / num;

  for (int x = 0; x < num; ++x) {
    float f = -(EARQ * MINBW) + ((hf + EARQ * MINBW)/exp(x*stepfactor/EARQ));
    printf("Filter @ %f\n", f);
    gts[x] = gen_gammatone(f, gt_lens+x);
  }
}

float *xcoeffs[NUMK][NUMK];

void xcorr(float **gt, int *gt_len) {
  for (int x = 0; x < NUMK; ++x) {
    for (int y = 0; y < NUMK; ++y) {
      float max = 0.0;
      xcoeffs[x][y] = (float *) malloc(sizeof(float)*gt_len[y]);
      for (int a = 0; a < gt_len[y]; ++a) {
	xcoeffs[x][y][a] = 0.0;
	for (int b = 0; b < gt_len[x] && a+b < gt_len[y]; ++b) {
	  xcoeffs[x][y][a] += gt[y][a+b]*gt[x][b];
	}
	if (xcoeffs[x][y][a] > max)
	  max = xcoeffs[x][y][a];
      }
      //if (max > 0.01) {
	neighborhoods[x].insert(y);
	neighborhoods[y].insert(x);
	//}
    }
  }
  
  for (int x = 0; x < NUMK; ++x) {
    if (!neighborhoods[x].empty()) {
      printf("%d: ", x);
      for (set<int>::const_iterator it = neighborhoods[x].begin();
	   it != neighborhoods[x].end(); ++it) {
	printf("%d ", *it);
      }
      printf("\n");
    }
  }
}

void load_sig(SNDFILE *isf, SF_INFO *sfi) {
  float *in_sig = (float *) malloc(sizeof(float)*sfi->channels*sfi->frames);
  sig_len = sfi->frames;
  sig = (float *) malloc(sizeof(float)*sig_len);
  orig_sig = (float *) malloc(sizeof(float)*sig_len);
  prev_sig = (float *) malloc(sizeof(float)*sig_len);

  sf_count_t total = sfi->channels*sfi->frames;
  sf_count_t offset = 0;
  sf_count_t cnt = 0;  
  sf_count_t readval = 10240;
  do {
    if ((total - offset)/2 < 10240)
      readval = (total-offset)/2;
    cnt = sf_readf_float(isf, in_sig + offset, readval);
    offset += sfi->channels*cnt;
  } while(offset < total);

  for (int x = 0; x < sig_len; ++x) {
    sig[x] = in_sig[2*x];
  }

  memcpy(orig_sig, sig, sizeof(float)*sig_len);

  free(in_sig);

  fprintf(stderr, "Read %d frames...\n", sig_len);
}

void dot_product(float *r, int rlen,
		 float *kernel, int klen,
		 float *c,
		 int start, int end,
		 int *t, float *coeff,
		 data_heap *h) {
  int y;
  float max = 0.0;
  if (start < 0)
    start = 0;
  if (end > rlen)
    end = rlen;
  float *curr_r = r+start;
  for (int x = start; x < end; ++x, ++curr_r) {
    c[x] = 0.0;
    float *val = curr_r;
    for (y = 0; (y < klen) && (x+y < rlen); ++y, ++val) {
      c[x] += (*val)*(kernel[y]);
    }
    h->push(c+x);
  }
  *coeff = 0.0;
  float *cmax = h->top();
  *coeff = *cmax;
  *t = cmax - c;
}

void update_dot_prod(float *c, float coeff, int offset,
		     int k1, int k1_len,
		     int k2, int k2_len,
		     int *t, float *cmax,
		     data_heap *h) {
  int y;

  for (y = offset - k1_len + 1; y < offset; ++y) {
    c[y] -= coeff*xcoeffs[k2][k1][offset-y];
    h->update(c+y);
  }

  for (; y < offset + k2_len && y < sig_len; ++y) {
    c[y] -= coeff*xcoeffs[k1][k2][y-offset];
    h->update(c+y);
  }

  *cmax = 0.0;
  float *max = h->top();
  *cmax = *max;
  *t = max - c;
}

int main(int argc, char *argv[]) {
  SF_INFO sfinfo;
  sfinfo.format = 0;
  SNDFILE* sf = sf_open(argv[1], SFM_READ, &sfinfo);

  fprintf(stderr,"%s:\n%lld\n%d\n%d\n%0x\n%d\n%d\n",
	  argv[1],
	  sfinfo.frames, sfinfo.samplerate,
	  sfinfo.channels, sfinfo.format,
	  sfinfo.sections, sfinfo.seekable);

  float *gt[NUMK];
  int gt_len[NUMK];
  fprintf(stderr, "Generating gammatones...\n");
  compute_gammatones(gt, gt_len, NUMK, 80.0, 44100.0);

  fprintf(stderr, "Computing xcorr...\n");
  xcorr(gt,gt_len);
  
  //sfinfo.frames = 44100;
  double s = sfinfo.frames / sfinfo.samplerate;

  load_sig(sf, &sfinfo);
  sf_close(sf);

  vector<float> coeffs;
  vector<int> ts;
  vector<int> ks;
  int indices[NUMK];
  int curr_ts[NUMK];
  float *c[NUMK];
  for (int x = 0; x < NUMK; ++x) {
    c[x] = (float *) malloc(sizeof(float)*sig_len);
    sig_heaps[x] = new data_heap(sig_len,
				 Cmp(),
				 ptr_map(c[x], new int[sig_len]));
  }

  fprintf(stderr, "Computing dot products...\n");
  for (int y = 0; y < NUMK; ++y) {
    indices[y] = y;
    dot_product(sig, sig_len, gt[y], gt_len[y], 
		c[y], 0, sig_len, curr_ts+y, curr_coeffs+y,
		sig_heaps[y]);
    fprintf(stderr, "%d ", y);
    fflush(stderr);
  }
  fprintf(stderr, "\n");
  for (int y = 0; y < NUMK; ++y)
    fprintf(stderr, "%f ", curr_coeffs[y]);
  fprintf(stderr, "\n");
  
  int ofile = open(argv[4], O_WRONLY|O_CREAT|O_TRUNC, 0666);

  float *itMax = max_element(curr_coeffs, curr_coeffs+NUMK);
  
  int best = itMax - curr_coeffs;
  
  int x = 0;
  while (x < 5000*s) {
    if (x % 100 == 0)
      fprintf(stderr, "%d %f\n", x, curr_coeffs[best]);

    ks.push_back(best);
    coeffs.push_back(curr_coeffs[best]);
    ts.push_back(curr_ts[best]);

    char kern = best;
    float coeff = curr_coeffs[best];
    int offset = curr_ts[best];

    write(ofile, &kern, sizeof(char));
    write(ofile, &coeff, sizeof(float));    
    write(ofile, &offset, sizeof(int));
    
    
    float *tmp = prev_sig;
    int z;
    prev_sig = sig;
    sig = tmp;

    for (z = 0; z < ts[x]; ++z)
      sig[z] = prev_sig[z];
    for (; z < ts[x]+gt_len[ks[x]] && z < sig_len; ++z)
      sig[z] = prev_sig[z] - coeffs[x]*gt[ks[x]][z-ts[x]];
    for(; z < sig_len; ++z)
      sig[z] = prev_sig[z];

    for (set<int>::const_iterator it = neighborhoods[best].begin();
	 it != neighborhoods[best].end(); ++it) {
      int val = *it;
      update_dot_prod(c[val], coeff, offset,
		      val, gt_len[val], best, gt_len[best],
		      curr_ts + val, curr_coeffs + val,
		      sig_heaps[val]);
    }

    itMax = max_element(curr_coeffs, curr_coeffs+NUMK);
    best = itMax - curr_coeffs;

    ++x;
  }

  close(ofile);

  float *of = (float *) calloc(sfinfo.channels*sfinfo.frames, sizeof(float));
  float *resid = (float *) malloc(sfinfo.channels*sfinfo.frames*sizeof(float));
  SF_INFO osfinfo;
  memcpy(&osfinfo, &sfinfo, sizeof(SF_INFO));
  SNDFILE* osf = sf_open(argv[2], SFM_WRITE, &osfinfo);
  memcpy(&osfinfo, &sfinfo, sizeof(SF_INFO));
  SNDFILE* rsf = sf_open(argv[3], SFM_WRITE, &osfinfo);
  for (unsigned int x = 0; x < ts.size(); ++x) {
    for (int y = 0; y < gt_len[ks[x]] && ts[x]+y < sig_len; ++y) {
      int tmp = 2*(ts[x]+y);
      of[tmp] += coeffs[x]*gt[ks[x]][y];
      of[tmp+1] = of[tmp];
    }
  }
  for (int x = 0; x < sig_len; ++x) {
    resid[2*x] = orig_sig[x] - of[2*x];
    resid[2*x+1] = resid[2*x];
  }
  sf_writef_float(osf, of, sig_len);
  sf_writef_float(rsf, resid, sig_len);
  sf_close(osf);
  sf_close(rsf);
  free(of);
  free(sig);

  return 0;
};
